# -*- coding: utf-8 -*-
from naie.datasets import get_data_reference
from naie.context import Context
import moxing as mox

import zipfile
import os
import time

import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
from torchvision import models

import numpy as np
import pandas as pd
from PIL import Image
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True


# ================================ 路径设定 ===============================
def create_the_dir(THE_DIR):
    if not os.path.exists(THE_DIR):
        os.makedirs(THE_DIR)


def create_these_dirs(THESE_DIRS):
    for THE_DIR in THESE_DIRS:
        create_the_dir(THE_DIR)


CACHE_DIR = r'./cahce/'
RAW_DIR = CACHE_DIR + r'/raw/'
LOAD_CHECKPOINT_DIR = CACHE_DIR + r'/prepared_checkpoints/'
SAVE_CHECKPOINT_DIR = CACHE_DIR + r'/produced_checkpoints/'
SAVE_TXT_DIR = CACHE_DIR + r'/logs/'
create_these_dirs([CACHE_DIR, 
                   RAW_DIR, 
                   LOAD_CHECKPOINT_DIR, 
                   SAVE_CHECKPOINT_DIR, 
                   SAVE_TXT_DIR])


# ============================ 训练超参数 =================================
NUM_CLASSES = 80

BATCH_SIZE = 64
NUM_WORKERS = 4
LR = 0.0002

# 训练轮数
MAX_EPOCH = 2
SAVE_EPOCH = 50

# 载入某一断点文件的路径
LOAD_PRETRAINED_PATH = LOAD_CHECKPOINT_DIR + 'resnet101-5d3b4d8f.pth'
LOAD_CHECKPOINT_PATH = None


# ================================== 导入云上文件的函数 ======================================
# 导入上传好的断点模型文件
def load_pretrained_checkpoints(dataset, dataset_entity, receive_dir):
    data_reference = get_data_reference(dataset, dataset_entity)
    for file_path in data_reference.get_files_paths():
        mox.file.copy(file_path, receive_dir + file_path.split('/')[-1])


# 导入在线数据集，并解压
def load_online_dataset(dataset, dataset_entity, receive_dir, password=None):
    data_reference = get_data_reference(dataset, dataset_entity)
    for file_path in data_reference.get_files_paths():
        mox.file.copy(file_path, receive_dir + file_path.split('/')[-1])

    filenames = os.listdir(receive_dir)
    for filename in filenames:
        if filename.endswith('.zip'):
            zip_file = zipfile.ZipFile(receive_dir + filename)
            zip_list = zip_file.namelist()
            for f in zip_list:
                if password is not None:
                    zip_file.extract(f, receive_dir, pwd=password.encode('utf-8'))
                else:
                    zip_file.extract(f, receive_dir)


if __name__ == '__main__':
    load_pretrained_checkpoints(dataset='Pretrained', dataset_entity='resnet', receive_dir=LOAD_CHECKPOINT_DIR)
    load_online_dataset(dataset='Contest', dataset_entity='ocean_new', receive_dir=RAW_DIR, password='Amphiprion')
